#!/usr/bin/env python3
"""
创建测试数据的脚本
"""
import sys
import os
sys.path.append(os.path.dirname(os.path.abspath(__file__)))

from sqlalchemy.orm import Session
from core.database import SessionLocal
from models.petition_record import PetitionRecord, RegionalAnalysisResult
from models.user import User
from models.analysis_task import AnalysisTask, TaskStatus, TaskType
from datetime import datetime

def create_test_data():
    """创建测试数据"""
    db = SessionLocal()
    try:
        # 获取admin用户
        admin_user = db.query(User).filter(User.username == "admin").first()
        if not admin_user:
            print("Admin user not found")
            return
        
        # 获取现有的信访记录
        existing_records = db.query(PetitionRecord).limit(3).all()
        if not existing_records:
            print("No existing petition records found")
            return
        
        # 创建一个分析任务
        analysis_task = AnalysisTask(
            name="测试分析任务",
            description="用于测试的分析任务",
            task_type=TaskType.RegionalAnalysis,
            config={},
            services=[],
            creator_id=admin_user.id,
            status=TaskStatus.COMPLETED
        )
        db.add(analysis_task)
        db.commit()
        
        # 为每个信访记录创建区域分析结果
        for i, record in enumerate(existing_records):
            analysis_result = RegionalAnalysisResult(
                petition_record_id=record.xh,  # 使用xh作为外键
                task_id=analysis_task.id,  # 使用刚创建的任务ID
                ccjbz=record.ccjbz,
                tsr=record.tsr,
                xfnr=record.xfnr,
                ts_gk=record.ts_gk,
                ts_xzq=record.ts_xzq,
                ts_qxdw=record.ts_qxdw,
                ts_dfnr=record.ts_dfnr,
                ts_bz=record.ts_bz,
                dzxx_tsdx="投诉对象",
                dzxx_tsdz=record.ts_dz,
                dzxx_xzq=record.ts_qxs,
                dzxx_jd=record.ts_jd,
                dzxx_lng=str(113.9 + i * 0.01),  # 经度
                dzxx_lat=str(22.5 + i * 0.01),   # 纬度
                cfxs_gl=str(85 + i * 5),  # 概率
            )
            db.add(analysis_result)
        
        db.commit()
        
        print(f"成功创建 {len(existing_records)} 条信访记录和分析结果")
        
    except Exception as e:
        print(f"创建测试数据失败: {e}")
        db.rollback()
    finally:
        db.close()

if __name__ == "__main__":
    create_test_data()